from typing import TypeVar

import numpy as np

from zkl_aiutils_datasets.basics import Dataset
from zkl_aiutils_datasets.processing import DelegatingDataset, SlicedDataset
from zkl_aiutils_datasets.wrapping import CountedDataset
from .random_gathering import RandomGatheringDataset

AnySample = TypeVar('AnySample')


class ShardedShuffledDataset(DelegatingDataset[AnySample]):
    def __init__(self, dataset: Dataset[AnySample], shards_num: int = 256):
        samples_n = len(dataset)
        shards_head = np.arange(shards_num) * (samples_n / shards_num)
        shards_head = np.asarray(shards_head, dtype=np.int64)
        shards_head = np.concat([[0], shards_head[1:][np.diff(shards_head) > 0]])
        shards_tail = np.concat([shards_head[1:], [samples_n]])
        shards_size = shards_tail - shards_head

        shards_dataset = [
            SlicedDataset(dataset, slice(shard_head, shard_tail))
            for shard_head, shard_tail in zip(shards_head, shards_tail)]

        dataset = RandomGatheringDataset(
            (dataset, weight)
            for dataset, weight in zip(shards_dataset, shards_size))

        dataset = CountedDataset(dataset, samples_n)

        super().__init__(dataset)
